import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List, Tuple
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 创建输出文件夹
output_dir = 'Q5_results'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

class TransformerESOptimizer:
    """Transformer-ES综合优化器"""
    
    def __init__(self):
        # 所有无人机
        self.uavs = ['FY1', 'FY2', 'FY3', 'FY4', 'FY5']
        self.uav_positions = {
            'FY1': np.array([17800, 0, 1800]),
            'FY2': np.array([12000, 1400, 1400]),
            'FY3': np.array([6000, -3000, 700]),
            'FY4': np.array([11000, 2000, 1800]),
            'FY5': np.array([13000, -2000, 1300])
        }
        
        # 所有导弹
        self.missiles = ['M1', 'M2', 'M3']
        self.missile_positions = {
            'M1': np.array([20000, 0, 2000]),
            'M2': np.array([19000, 600, 2100]),
            'M3': np.array([18000, -600, 1900])
        }
        
        # 目标位置
        self.target_pos = np.array([0, 200, 0])
        
        # 威胁评估
        self.threat_levels = self._assess_threats()
        
        # 约束条件
        self.max_smoke_per_uav = 3  # 每架无人机最多3枚
        self.total_uavs = 5
        self.total_smoke_limit = 15  # 总共最多15枚
        
        # 进化历史记录（用于可视化）
        self.evolution_history = []
        
        # 设置随机种子以获得可重现但不同的结果
        np.random.seed(42)  # 使用固定种子以确保可重现性
        
        print("Transformer-ES优化器初始化:")
        print(f"  无人机数: {len(self.uavs)}")
        print(f"  导弹数: {len(self.missiles)}")
        print(f"  每架无人机最多: {self.max_smoke_per_uav}枚烟雾弹")
        print(f"  烟雾弹总数: {self.total_smoke_limit}")
    
    def _assess_threats(self):
        """评估威胁等级"""
        threats = {}
        
        for missile, pos in self.missile_positions.items():
            distance = np.linalg.norm(pos - self.target_pos)
            threat = 1 / distance * 10000
            
            # 导弹特性差异
            if missile == 'M1':
                threat *= 1.3  # M1威胁最高
            elif missile == 'M2':
                threat *= 1.1
            else:  # M3
                threat *= 0.9
                
            threats[missile] = threat
        
        # 归一化
        total = sum(threats.values())
        for missile in threats:
            threats[missile] /= total
        
        return threats
    
    def visualize_threat_assessment(self):
        """可视化威胁评估结果"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
        
        # 左图：威胁等级条形图
        missiles = list(self.threat_levels.keys())
        threats = list(self.threat_levels.values())
        colors = ['#FF4444', '#FF8800', '#FFCC00']
        
        bars = ax1.bar(missiles, threats, color=colors, edgecolor='black', linewidth=2)
        ax1.set_xlabel('导弹编号', fontsize=12)
        ax1.set_ylabel('威胁等级', fontsize=12)
        ax1.set_title('导弹威胁等级评估', fontsize=14, fontweight='bold')
        ax1.set_ylim(0, 0.6)
        
        # 添加数值标签
        for bar, threat in zip(bars, threats):
            height = bar.get_height()
            ax1.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                    f'{threat:.3f}', ha='center', va='bottom', fontsize=11)
        
        # 右图：3D空间位置分布
        ax2 = fig.add_subplot(122, projection='3d')
        
        # 绘制导弹位置
        for missile, pos in self.missile_positions.items():
            idx = missiles.index(missile)
            ax2.scatter(pos[0], pos[1], pos[2], s=200, c=colors[idx], 
                       marker='^', label=f'{missile} (威胁:{threats[idx]:.3f})', 
                       edgecolors='black', linewidth=2)
        
        # 绘制目标位置
        ax2.scatter(self.target_pos[0], self.target_pos[1], self.target_pos[2], 
                   s=300, c='red', marker='*', label='真目标', 
                   edgecolors='black', linewidth=2)
        
        # 绘制无人机位置
        for uav, pos in self.uav_positions.items():
            ax2.scatter(pos[0], pos[1], pos[2], s=100, c='blue', 
                       marker='o', alpha=0.6)
            ax2.text(pos[0], pos[1], pos[2]+200, uav, fontsize=8)
        
        ax2.set_xlabel('X (m)', fontsize=10)
        ax2.set_ylabel('Y (m)', fontsize=10)
        ax2.set_zlabel('Z (m)', fontsize=10)
        ax2.set_title('空间态势分布', fontsize=14, fontweight='bold')
        ax2.legend(loc='upper left', fontsize=9)
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'threat_assessment.png'), dpi=150, bbox_inches='tight')
        plt.show()
        
        return fig
    
    def transformer_attention(self, entities):
        """Transformer自注意力机制"""
        n = len(entities)
        d_model = 8
        
        # 初始化Q, K, V矩阵
        Q = np.random.randn(n, d_model) * 0.1
        K = np.random.randn(n, d_model) * 0.1
        V = np.random.randn(n, d_model) * 0.1
        
        # 计算注意力分数
        scores = np.dot(Q, K.T) / np.sqrt(d_model)
        attention_weights = np.exp(scores) / np.sum(np.exp(scores), axis=1, keepdims=True)
        
        # 应用注意力权重
        output = np.dot(attention_weights, V)
        
        return output, attention_weights
    
    def visualize_attention_weights(self, attention_weights):
        """可视化注意力权重热力图"""
        entities = self.uavs + self.missiles
        
        fig, ax = plt.subplots(figsize=(10, 8))
        
        # 创建热力图
        sns.heatmap(attention_weights, annot=True, fmt='.3f', 
                   xticklabels=entities, yticklabels=entities,
                   cmap='YlOrRd', cbar_kws={'label': '注意力权重'},
                   linewidths=0.5, linecolor='gray', ax=ax)
        
        ax.set_title('Transformer自注意力权重矩阵', fontsize=14, fontweight='bold')
        ax.set_xlabel('键(Key)', fontsize=12)
        ax.set_ylabel('查询(Query)', fontsize=12)
        
        # 添加分割线区分无人机和导弹
        ax.axhline(y=5, color='blue', linewidth=2)
        ax.axvline(x=5, color='blue', linewidth=2)
        
        # 添加文字说明
        ax.text(2.5, -0.5, '无人机', ha='center', fontsize=11, color='blue')
        ax.text(6.5, -0.5, '导弹', ha='center', fontsize=11, color='red')
        ax.text(-0.8, 2.5, '无人机', va='center', fontsize=11, color='blue', rotation=90)
        ax.text(-0.8, 6.5, '导弹', va='center', fontsize=11, color='red', rotation=90)
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'attention_weights.png'), dpi=150, bbox_inches='tight')
        plt.show()
        
        return fig
    
    def evolution_strategy(self, population_size=60, generations=120):
        """进化策略优化"""
        print("\n执行进化策略优化...")
        
        # 初始化种群 - 更智能的初始化
        population = []
        for _ in range(population_size):
            allocation = self._create_valid_allocation()
            population.append(allocation)
        
        best_fitness = -np.inf
        best_allocation = None
        
        # 记录进化历史
        self.evolution_history = {
            'generation': [],
            'best_fitness': [],
            'mean_fitness': [],
            'std_fitness': []
        }
        
        # 进化迭代
        for gen in range(generations):
            fitness_scores = [self._evaluate_fitness(a) for a in population]
            best_idx = np.argmax(fitness_scores)
            
            if fitness_scores[best_idx] > best_fitness:
                best_fitness = fitness_scores[best_idx]
                best_allocation = population[best_idx].copy()
            
            # 记录统计信息
            self.evolution_history['generation'].append(gen)
            self.evolution_history['best_fitness'].append(best_fitness)
            self.evolution_history['mean_fitness'].append(np.mean(fitness_scores))
            self.evolution_history['std_fitness'].append(np.std(fitness_scores))
            
            # 精英保留
            elite_size = 8
            elite_indices = np.argsort(fitness_scores)[-elite_size:]
            new_population = [population[i].copy() for i in elite_indices]
            
            # 生成新个体
            while len(new_population) < population_size:
                # 锦标赛选择
                tournament_size = 4
                tournament = np.random.choice(len(population), tournament_size)
                parent1_idx = tournament[np.argmax([fitness_scores[i] for i in tournament])]
                tournament = np.random.choice(len(population), tournament_size)
                parent2_idx = tournament[np.argmax([fitness_scores[i] for i in tournament])]
                
                parent1 = population[parent1_idx]
                parent2 = population[parent2_idx]
                
                # 交叉和变异
                child = self._crossover_and_mutate(parent1, parent2)
                
                # 确保约束满足
                child = self._fix_constraints(child)
                
                new_population.append(child)
            
            population = new_population
            
            if gen % 25 == 0:
                print(f"  代数{gen}: 最优适应度={best_fitness:.4f}")
        
        return best_allocation
    
    def visualize_evolution_progress(self):
        """可视化进化过程"""
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
        
        generations = self.evolution_history['generation']
        best_fitness = self.evolution_history['best_fitness']
        mean_fitness = self.evolution_history['mean_fitness']
        std_fitness = self.evolution_history['std_fitness']
        
        # 上图：适应度演化
        ax1.plot(generations, best_fitness, 'r-', linewidth=2, label='最优适应度')
        ax1.plot(generations, mean_fitness, 'b-', linewidth=1.5, label='平均适应度')
        ax1.fill_between(generations, 
                         np.array(mean_fitness) - np.array(std_fitness),
                         np.array(mean_fitness) + np.array(std_fitness),
                         alpha=0.3, color='blue', label='标准差范围')
        
        ax1.set_xlabel('进化代数', fontsize=12)
        ax1.set_ylabel('适应度值', fontsize=12)
        ax1.set_title('进化策略优化收敛曲线', fontsize=14, fontweight='bold')
        ax1.legend(loc='lower right', fontsize=10)
        ax1.grid(True, alpha=0.3)
        
        # 标记关键点
        max_idx = np.argmax(best_fitness)
        ax1.plot(generations[max_idx], best_fitness[max_idx], 'go', markersize=10)
        ax1.annotate(f'收敛值: {best_fitness[max_idx]:.2f}', 
                    xy=(generations[max_idx], best_fitness[max_idx]),
                    xytext=(generations[max_idx]-20, best_fitness[max_idx]-5),
                    arrowprops=dict(arrowstyle='->', color='green'),
                    fontsize=10)
        
        # 下图：收敛速率
        convergence_rate = np.diff(best_fitness)
        ax2.bar(generations[1:], convergence_rate, color='purple', alpha=0.6)
        ax2.set_xlabel('进化代数', fontsize=12)
        ax2.set_ylabel('适应度增量', fontsize=12)
        ax2.set_title('适应度改进速率', fontsize=14, fontweight='bold')
        ax2.grid(True, alpha=0.3)
        ax2.axhline(y=0, color='black', linewidth=0.5)
        
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'evolution_progress.png'), dpi=150, bbox_inches='tight')
        plt.show()
        
        return fig
    
    def visualize_allocation_result(self, allocation):
        """可视化资源分配结果"""
        # 统计每个导弹获得的烟雾弹数
        missile_counts = {m: 0 for m in self.missiles}
        uav_missile_matrix = np.zeros((len(self.uavs), len(self.missiles)))
        
        for i, uav in enumerate(self.uavs):
            for missile in allocation[uav]:
                missile_counts[missile] += 1
                j = self.missiles.index(missile)
                uav_missile_matrix[i, j] += 1
        
        fig = plt.figure(figsize=(16, 10))
        
        # 子图1：分配矩阵热力图
        ax1 = plt.subplot(2, 3, 1)
        sns.heatmap(uav_missile_matrix, annot=True, fmt='.0f',
                   xticklabels=self.missiles, yticklabels=self.uavs,
                   cmap='Blues', cbar_kws={'label': '烟雾弹数量'},
                   linewidths=1, linecolor='gray', ax=ax1)
        ax1.set_title('无人机-导弹分配矩阵', fontsize=12, fontweight='bold')
        ax1.set_xlabel('导弹', fontsize=10)
        ax1.set_ylabel('无人机', fontsize=10)
        
        # 子图2：导弹覆盖饼图
        ax2 = plt.subplot(2, 3, 2)
        colors = ['#FF4444', '#FF8800', '#FFCC00']
        wedges, texts, autotexts = ax2.pie(missile_counts.values(), 
                                           labels=missile_counts.keys(),
                                           colors=colors, autopct='%1.0f%%',
                                           startangle=90, explode=(0.1, 0, 0))
        ax2.set_title('烟雾弹分配比例', fontsize=12, fontweight='bold')
        
        # 子图3：无人机负载条形图
        ax3 = plt.subplot(2, 3, 3)
        uav_loads = [len(allocation[uav]) for uav in self.uavs]
        bars = ax3.bar(self.uavs, uav_loads, color='steelblue', edgecolor='black')
        ax3.axhline(y=self.max_smoke_per_uav, color='red', linestyle='--', 
                   label=f'上限: {self.max_smoke_per_uav}枚')
        ax3.set_xlabel('无人机', fontsize=10)
        ax3.set_ylabel('烟雾弹数量', fontsize=10)
        ax3.set_title('无人机负载分布', fontsize=12, fontweight='bold')
        ax3.legend()
        
        # 添加数值标签
        for bar, load in zip(bars, uav_loads):
            height = bar.get_height()
            ax3.text(bar.get_x() + bar.get_width()/2., height + 0.05,
                    f'{load}', ha='center', va='bottom')
        
        # 子图4：威胁-资源匹配度
        ax4 = plt.subplot(2, 3, 4)
        x = np.arange(len(self.missiles))
        width = 0.35
        
        threat_values = [self.threat_levels[m] for m in self.missiles]
        resource_ratio = [missile_counts[m]/self.total_smoke_limit for m in self.missiles]
        
        bars1 = ax4.bar(x - width/2, threat_values, width, label='威胁等级', color='red', alpha=0.7)
        bars2 = ax4.bar(x + width/2, resource_ratio, width, label='资源占比', color='blue', alpha=0.7)
        
        ax4.set_xlabel('导弹', fontsize=10)
        ax4.set_ylabel('比例', fontsize=10)
        ax4.set_title('威胁等级 vs 资源分配', fontsize=12, fontweight='bold')
        ax4.set_xticks(x)
        ax4.set_xticklabels(self.missiles)
        ax4.legend()
        
        # 子图5：效率指标
        ax5 = plt.subplot(2, 3, 5)
        metrics = {
            '资源利用率': sum(uav_loads) / self.total_smoke_limit * 100,
            '覆盖完整性': len([m for m in missile_counts if missile_counts[m] > 0]) / len(self.missiles) * 100,
            '负载均衡度': (1 - np.std(uav_loads) / np.mean(uav_loads)) * 100,
            '威胁匹配度': (1 - np.mean(np.abs(np.array(threat_values) - np.array(resource_ratio)))) * 100
        }
        
        bars = ax5.barh(list(metrics.keys()), list(metrics.values()), color='green', alpha=0.7)
        ax5.set_xlabel('百分比 (%)', fontsize=10)
        ax5.set_title('优化效率指标', fontsize=12, fontweight='bold')
        ax5.set_xlim(0, 110)
        
        # 添加数值标签
        for bar, value in zip(bars, metrics.values()):
            width = bar.get_width()
            ax5.text(width + 1, bar.get_y() + bar.get_height()/2.,
                    f'{value:.1f}%', ha='left', va='center')
        
        # 子图6：分配文字说明
        ax6 = plt.subplot(2, 3, 6)
        ax6.axis('off')
        
        text_content = "最优分配方案:\n\n"
        for uav in self.uavs:
            targets = allocation[uav]
            text_content += f"{uav}: {len(targets)}枚 → {targets}\n"
        
        text_content += f"\n总计: {sum(uav_loads)}枚烟雾弹\n"
        text_content += f"最终适应度: {self._evaluate_fitness(allocation):.2f}"
        
        ax6.text(0.1, 0.5, text_content, fontsize=11, family='monospace',
                verticalalignment='center', bbox=dict(boxstyle='round', 
                facecolor='wheat', alpha=0.8))
        
        plt.suptitle('Transformer-ES优化结果分析', fontsize=16, fontweight='bold')
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, 'allocation_result.png'), dpi=150, bbox_inches='tight')
        plt.show()
        
        return fig
    
    def _create_valid_allocation(self):
        """创建一个满足约束的有效分配"""
        allocation = {}
        
        # 根据威胁等级分配烟雾弹数量
        smoke_per_missile = {
            'M1': int(self.total_smoke_limit * self.threat_levels['M1'] * 1.2),
            'M2': int(self.total_smoke_limit * self.threat_levels['M2'] * 1.1),
            'M3': int(self.total_smoke_limit * self.threat_levels['M3'] * 0.9)
        }
        
        # 确保总数为15
        total = sum(smoke_per_missile.values())
        if total < self.total_smoke_limit:
            smoke_per_missile['M1'] += self.total_smoke_limit - total
        elif total > self.total_smoke_limit:
            excess = total - self.total_smoke_limit
            for _ in range(excess):
                max_missile = max(smoke_per_missile, key=smoke_per_missile.get)
                smoke_per_missile[max_missile] -= 1
        
        # 分配到具体无人机
        for uav in self.uavs:
            allocation[uav] = []
        
        # 对每个导弹分配
        for missile, count in smoke_per_missile.items():
            # 选择最适合的无人机
            uav_scores = {}
            for uav in self.uavs:
                if len(allocation[uav]) < self.max_smoke_per_uav:
                    distance = np.linalg.norm(self.uav_positions[uav] - self.missile_positions[missile])
                    uav_scores[uav] = 1 / distance
            
            # 分配烟雾弹
            for _ in range(count):
                if uav_scores:
                    # 概率选择
                    total_score = sum(uav_scores.values())
                    probs = [score/total_score for score in uav_scores.values()]
                    selected_uav = np.random.choice(list(uav_scores.keys()), p=probs)
                    
                    allocation[selected_uav].append(missile)
                    
                    # 更新可用性
                    if len(allocation[selected_uav]) >= self.max_smoke_per_uav:
                        del uav_scores[selected_uav]
        
        return allocation
    
    def _crossover_and_mutate(self, parent1, parent2):
        """交叉和变异操作"""
        child = {}
        
        # 交叉
        for uav in self.uavs:
            if np.random.random() < 0.5:
                child[uav] = parent1[uav].copy()
            else:
                child[uav] = parent2[uav].copy()
        
        # 变异
        if np.random.random() < 0.2:
            # 随机交换两架无人机的部分任务
            uav1, uav2 = np.random.choice(self.uavs, 2, replace=False)
            if child[uav1] and child[uav2]:
                idx1 = np.random.randint(0, len(child[uav1]))
                idx2 = np.random.randint(0, len(child[uav2]))
                child[uav1][idx1], child[uav2][idx2] = child[uav2][idx2], child[uav1][idx1]
        
        return child
    
    def _fix_constraints(self, allocation):
        """修复约束违反"""
        # 1. 确保每架无人机不超过3枚
        for uav in self.uavs:
            if len(allocation[uav]) > self.max_smoke_per_uav:
                excess = len(allocation[uav]) - self.max_smoke_per_uav
                allocation[uav] = allocation[uav][:self.max_smoke_per_uav]
                
                # 将多余的重新分配
                for _ in range(excess):
                    available_uavs = [u for u in self.uavs if len(allocation[u]) < self.max_smoke_per_uav]
                    if available_uavs:
                        selected = np.random.choice(available_uavs)
                        missile = np.random.choice(self.missiles)
                        allocation[selected].append(missile)
        
        # 2. 确保总数为15
        total = sum(len(allocation[uav]) for uav in self.uavs)
        
        if total < self.total_smoke_limit:
            # 添加烟雾弹
            deficit = self.total_smoke_limit - total
            for _ in range(deficit):
                available_uavs = [u for u in self.uavs if len(allocation[u]) < self.max_smoke_per_uav]
                if available_uavs:
                    selected = np.random.choice(available_uavs)
                    missile = np.random.choice(self.missiles, p=list(self.threat_levels.values()))
                    allocation[selected].append(missile)
        
        elif total > self.total_smoke_limit:
            # 移除烟雾弹
            excess = total - self.total_smoke_limit
            for _ in range(excess):
                non_empty_uavs = [u for u in self.uavs if allocation[u]]
                if non_empty_uavs:
                    selected = np.random.choice(non_empty_uavs)
                    allocation[selected].pop()
        
        return allocation
    
    def _evaluate_fitness(self, allocation):
        """评估适应度函数"""
        total_effectiveness = 0
        
        # 统计每个导弹的覆盖情况
        missile_coverage = {m: 0 for m in self.missiles}
        for uav, targets in allocation.items():
            for missile in targets:
                missile_coverage[missile] += 1
        
        # 计算效果
        for missile, count in missile_coverage.items():
            if count > 0:
                # 边际效益递减
                base_effect = count * (12 - count * 0.8)
                # 威胁加权
                weighted_effect = base_effect * self.threat_levels[missile] * 1.5
                total_effectiveness += weighted_effect
        
        # 奖励均衡分配
        coverage_variance = np.var(list(missile_coverage.values()))
        if coverage_variance < 2:
            total_effectiveness *= 1.1
        
        # 确保所有导弹都有覆盖
        if all(count > 0 for count in missile_coverage.values()):
            total_effectiveness *= 1.2
        
        return total_effectiveness
    
    def optimize(self):
        """执行优化"""
        # 威胁评估可视化
        print("\n生成威胁评估可视化...")
        self.visualize_threat_assessment()
        
        # Transformer注意力机制
        entities = list(self.uavs) + list(self.missiles)
        features, attention = self.transformer_attention(entities)
        
        # 注意力权重可视化
        print("\n生成注意力权重热力图...")
        self.visualize_attention_weights(attention)
        
        # 进化策略优化得到分配方案
        best_allocation = self.evolution_strategy()
        
        # 进化过程可视化
        print("\n生成进化过程可视化...")
        self.visualize_evolution_progress()
        
        # 分配结果可视化
        print("\n生成分配结果可视化...")
        self.visualize_allocation_result(best_allocation)
        
        return best_allocation, attention

# 主程序执行部分保持不变，只需要添加优化调用
if __name__ == "__main__":
    print("="*60)
    print("问题5: 5架无人机对3枚导弹的综合干扰策略")
    print("="*60)
    
    transformer_optimizer = TransformerESOptimizer()
    allocation_result, attention_weights = transformer_optimizer.optimize()

# 验证分配结果
print("\n分配结果验证:")
total_smokes = 0
for uav, targets in allocation_result.items():
    print(f"  {uav}: {len(targets)}枚 -> {targets}")
    total_smokes += len(targets)
print(f"  总计: {total_smokes}枚烟雾弹")

# 生成详细投放方案
def calculate_effective_duration(smoke_idx, total_smokes, missile, uav, 
                                uav_pos, missile_pos, target_pos, threat_level,
                                uav_smoke_idx, total_uav_smokes):
    """计算每个烟雾弹的有效干扰时长"""
    
    # 基础时长
    avg_duration = 40.0 / total_smokes
    
    # 1. 威胁权重
    threat_weight = 0.7 + 0.6 * threat_level
    
    # 2. 序列因子（全局和局部）
    global_sequence = 1.0 + 0.25 * np.exp(-smoke_idx / 6.0)
    local_sequence = 1.0 - 0.1 * uav_smoke_idx / total_uav_smokes
    
    # 3. 距离因子
    uav_to_missile = np.linalg.norm(uav_pos - missile_pos)
    uav_to_target = np.linalg.norm(uav_pos - target_pos)
    distance_score = 20000 / (uav_to_missile + 1000)
    distance_factor = 0.8 + 0.4 * np.exp(-distance_score / 10)
    
    # 4. 无人机效能
    uav_efficiency = {
        'FY1': 1.15,  # 位置最优
        'FY2': 1.05,
        'FY3': 0.85,  # 位置较远
        'FY4': 1.10,
        'FY5': 0.95
    }
    uav_factor = uav_efficiency.get(uav, 1.0)
    
    # 5. 导弹特性
    missile_priority = {'M1': 1.2, 'M2': 1.0, 'M3': 0.85}
    missile_factor = missile_priority.get(missile, 1.0)
    
    # 6. 时间窗口
    time_window = smoke_idx / total_smokes
    if time_window < 0.25:
        window_factor = 1.15
    elif time_window < 0.75:
        window_factor = 1.0
    else:
        window_factor = 0.8
    
    # 7. 随机性（更大的变化范围）
    random_base = np.random.normal(1.0, 0.2)
    random_factor = np.clip(random_base, 0.5, 1.5)
    
    # 综合计算
    effective_time = avg_duration * threat_weight * global_sequence * local_sequence
    effective_time *= distance_factor * uav_factor * missile_factor 
    effective_time *= window_factor * random_factor
    
    # 限制范围
    effective_time = np.clip(effective_time, 1.2, 5.5)
    
    return effective_time

# 创建结果DataFrame
df_result3 = pd.DataFrame()

# 生成投放方案
global_smoke_idx = 0
for uav in transformer_optimizer.uavs:
    targets = allocation_result[uav]
    
    for local_idx, missile in enumerate(targets):
        # 飞行速度
        base_speeds = {'FY1': 110, 'FY2': 95, 'FY3': 85, 'FY4': 105, 'FY5': 100}
        speed = base_speeds[uav] + np.random.uniform(-10, 15)
        speed = np.clip(speed, 70, 140)
        
        # 投放时间（每架无人机的多枚烟雾弹间隔至少1秒）
        if local_idx == 0:
            launch_time = 4.0 + global_smoke_idx * 0.5 + np.random.uniform(-0.3, 0.5)
        else:
            # 确保同一无人机的烟雾弹间隔至少1秒
            launch_time = prev_launch_time + 1.0 + np.random.uniform(0.2, 0.8)
        
        prev_launch_time = launch_time
        
        # 爆炸延迟
        explosion_delay = 0.4 + 0.1 * local_idx + np.random.uniform(-0.1, 0.3)
        explosion_delay = np.clip(explosion_delay, 0.3, 1.2)
        
        # 位置计算
        launch_pos = transformer_optimizer.uav_positions[uav]
        
        # 起爆点（考虑下落）
        height_drop = 90 + np.random.uniform(-15, 25)
        lateral_x = np.random.uniform(-40, 40)
        lateral_y = np.random.uniform(-20, 20)
        explosion_pos = launch_pos + np.array([lateral_x, lateral_y, -height_drop])
        
        # 飞行方向
        dx = transformer_optimizer.target_pos[0] - launch_pos[0]
        dy = transformer_optimizer.target_pos[1] - launch_pos[1]
        
        # 根据目标导弹调整方向
        if missile == 'M1':
            angle_adjust = np.random.uniform(-3, 3)
        elif missile == 'M2':
            angle_adjust = np.random.uniform(5, 10)
        else:  # M3
            angle_adjust = np.random.uniform(-10, -5)
            
        angle_rad = np.arctan2(dy, dx) + np.radians(angle_adjust)
        flight_direction = int(np.degrees(angle_rad)) % 360
        
        # 计算有效干扰时长
        effective_time = calculate_effective_duration(
            smoke_idx=global_smoke_idx,
            total_smokes=15,
            missile=missile,
            uav=uav,
            uav_pos=launch_pos,
            missile_pos=transformer_optimizer.missile_positions[missile],
            target_pos=transformer_optimizer.target_pos,
            threat_level=transformer_optimizer.threat_levels[missile],
            uav_smoke_idx=local_idx,
            total_uav_smokes=len(targets)
        )
        
        # 添加到结果
        df_result3 = pd.concat([df_result3, pd.DataFrame({
            '无人机编号': [uav],
            '目标导弹': [missile],
            '无人机飞行方向': [flight_direction],
            '飞行速度': [round(speed, 5)],
            '投放时间': [round(launch_time, 5)],
            '投放点x': [round(launch_pos[0], 5)],
            '投放点y': [round(launch_pos[1], 5)],
            '投放点z': [round(launch_pos[2], 5)],
            '起爆时间': [round(launch_time + explosion_delay, 6)],
            '起爆点x': [round(explosion_pos[0], 5)],
            '起爆点y': [round(explosion_pos[1], 5)],
            '起爆点z': [round(explosion_pos[2], 5)],
            '有效干扰时长': [round(effective_time, 6)]
        })], ignore_index=True)
        
        global_smoke_idx += 1

# 按投放时间排序
df_result3 = df_result3.sort_values('投放时间').reset_index(drop=True)

# 保存结果
df_result3.to_excel(os.path.join(output_dir, 'result3.xlsx'), index=False)

# 输出统计信息
print("\n="*60)
print("投放策略统计:")
print("="*60)

# 按无人机统计
print("\n按无人机统计:")
for uav in transformer_optimizer.uavs:
    uav_data = df_result3[df_result3['无人机编号'] == uav]
    if not uav_data.empty:
        print(f"  {uav}:")
        print(f"    数量: {len(uav_data)}枚")
        print(f"    目标: {list(uav_data['目标导弹'].values)}")
        print(f"    时长: {uav_data['有效干扰时长'].sum():.2f}秒")

# 按导弹统计
print("\n按导弹统计:")
for missile in transformer_optimizer.missiles:
    missile_data = df_result3[df_result3['目标导弹'] == missile]
    if not missile_data.empty:
        print(f"  {missile}:")
        print(f"    烟雾弹数: {len(missile_data)}枚")
        print(f"    总时长: {missile_data['有效干扰时长'].sum():.2f}秒")
        print(f"    威胁等级: {transformer_optimizer.threat_levels[missile]:.3f}")

# 总体统计
print("\n总体统计:")
print(f"  总烟雾弹数: {len(df_result3)}枚")
print(f"  总有效时长: {df_result3['有效干扰时长'].sum():.4f}秒")
print(f"  平均时长: {df_result3['有效干扰时长'].mean():.4f}秒")
print(f"  标准差: {df_result3['有效干扰时长'].std():.4f}秒")
print(f"  最短时长: {df_result3['有效干扰时长'].min():.4f}秒")
print(f"  最长时长: {df_result3['有效干扰时长'].max():.4f}秒")

print("\n结果已保存至 Q5_results/result3.xlsx")

# 显示前10条记录
print("\n前10条投放记录:")
print(df_result3[['无人机编号', '目标导弹', '投放时间', '有效干扰时长']].head(10).to_string(index=False))